import numpy as np
import pandas as pd
from flask import Blueprint, jsonify, render_template, request
from sqlalchemy import func
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LinearRegression

from applications.extensions import db
from applications.common.utils.rights import authorize
from applications.models.bigScreenData import PlantingArea, AgriculturalYield, AffectedArea
from sklearn.metrics import r2_score, mean_squared_error, mean_absolute_error

bp = Blueprint('module', __name__, url_prefix='/module')

@authorize("system:user:main")
def prepare_data():
    models = [AgriculturalYield, AffectedArea, PlantingArea]
    # yield  产量   affected 受灾  planting 种植
    keys = ["yield", "affected", "planting"]
    yearly_data = []

    # 数据获取逻辑
    for model, key in zip(models, keys):
        results = db.session.query(
            model.year,
            func.coalesce(func.sum(model.value), 0).label('total')
        ).filter(
            model.value.isnot(None),
            model.value != 0.0
        ).group_by(model.year).all()

        for year, total in results:
            entry = next((item for item in yearly_data if item["year"] == year), None)
            if not entry:
                entry = {"year": year, "yield": 0.0, "affected": 0.0, "planting": 0.0}
                yearly_data.append(entry)
            if float(total) > 0:
                entry[key] = float(total)

    yearly_data = [
        item for item in yearly_data
        if any(item[key] > 0 for key in ["yield", "affected", "planting"])
    ]
    yearly_data.sort(key=lambda x: x["year"])

    df = pd.DataFrame(yearly_data)
    return df


@bp.post('/train')
@authorize("system:user:main")
def train_model():
    predict_years = request.get_json().get('predict_years')
    # 获取数据
    df = prepare_data()

    # 过滤掉任何属性为0的记录
    df = df[(df['yield'] > 0) & (df['affected'] > 0) & (df['planting'] > 0)]

    if df.empty:
        return jsonify({"success": False, "message": "没有足够的有效数据进行训练"})

    # 准备特征和目标变量
    X = df[['affected', 'planting']].values
    y = df['yield'].values

    # 标准化特征
    scaler = StandardScaler()
    X_scaled = scaler.fit_transform(X)

    # 训练线性回归模型
    model = LinearRegression()
    model.fit(X_scaled, y)

    # 预测未来5年的产量
    last_year = int(df['year'].max())
    future_years = [str(last_year + i) for i in range(1, int(predict_years) + 1)]

    # 使用最后一年的受灾面积和种植面积作为基准，预测未来产量
    last_data = df[df['year'] == str(last_year)]
    base_affected = last_data['affected'].values[0]
    base_planting = last_data['planting'].values[0]

    # 假设未来5年受灾面积每年减少2%，种植面积每年增加affected%
    future_predictions = []
    for i, year in enumerate(future_years):
        affected = base_affected * (0.98 ** (i + 1))  # 受灾面积每年减少2%
        planting = base_planting * (1.01 ** (i + 1))  # 种植面积每年增加1%
        # 标准化特征
        X_future = scaler.transform([[affected, planting]])

        # 预测产量
        predicted_yield = model.predict(X_future)[0]

        future_predictions.append({
            "year": int(year),
            "yield": round(predicted_yield, 2),
            "affected": round(affected, 2),
            "planting": round(planting, 2)
        })

    # 返回预测结果
    return jsonify({
        "success": True,
        "data": {
            "model_info": {
                "coefficients": dict(zip(['受灾面积', '种植面积'], model.coef_.tolist())),
                "intercept": float(model.intercept_),
                "metrics": {
                    "r2_score": round(r2_score(y, model.predict(X_scaled)), 4),
                    "mse": round(mean_squared_error(y, model.predict(X_scaled)), 2),
                    "mae": round(mean_absolute_error(y, model.predict(X_scaled)), 2),
                    "feature_importance": dict(zip(['受灾面积', '种植面积'], model.coef_.tolist()))
                }
            },
            "training_data": df.to_dict(orient="records"),
            "predictions": future_predictions
        }
    })


@bp.get('/index')
@authorize("system:user:main")
def index():
    return render_template('module/index.html')
